package org.zh.client.config;

import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Mono;

import java.util.List;

import static org.zh.client.config.GlobalStatic.GLOBAL_TRANSACTION;
import static org.zh.client.config.GlobalStatic.GLOBAL_ThreadLocal;

public class ClientWebFilter implements WebFilter {
    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
        List<String> strings = exchange.getRequest().getHeaders().get(GLOBAL_TRANSACTION);
        if (strings != null && !strings.isEmpty()) {
            GLOBAL_ThreadLocal.set(strings.getFirst());
            return chain.filter(exchange);
        }
        return chain.filter(exchange).doFinally(s -> {
            if (GLOBAL_ThreadLocal.get() != null) {
                GLOBAL_ThreadLocal.remove();
            }
        });
    }
}
